#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
CSV数据处理工具模块

提供CSV文件的读取、写入、移动和格式化功能
"""

import os
import shutil
import csv
import math
from datetime import datetime


class CSVUtils:
    """CSV处理工具类"""
    
    @staticmethod
    def create_simion_csv_with_extended_data(splats, frequency, target_dir, emittance_data=None, extended_params=None, prefix="simion"):
        """
        创建包含发射度信息和扩展参数的SIMION Splat数据CSV文件
        
        参数:
            splats: Splat数据列表 [(x, y, z, mass, vx, vy, vz, ke, t, charge), ...]
            frequency: 频率值
            target_dir: 目标目录
            emittance_data: 发射度数据字典
            extended_params: 扩展参数字典
            prefix: 文件名前缀
            
        返回:
            str: 创建的文件路径，如果失败返回None
        """
        try:
            # 确保目标目录存在
            os.makedirs(target_dir, exist_ok=True)
            
            # 生成包含所有信息的文件名（英文短标签，使用等号）
            if emittance_data and extended_params:
                filename = (
                    f"freq={frequency:.0f}Hz_"
                    f"rms={emittance_data.get('rms_x', 0):.2e}_"
                    f"geom={emittance_data.get('geometric_x', 0):.2e}_"
                    f"norm={emittance_data.get('normalized_x', 0):.2e}_"
                    f"hull={extended_params.get('convex_hull_area', 0):.2e}.csv"
                )
            else:
                filename = f"freq={frequency:.0f}Hz.csv"
            target_path = os.path.join(target_dir, filename)
            
            # 准备扩展的headers
            headers = ['x', 'y', 'z', 'mass', 'vx', 'vy', 'vz', 'ke', 't', 'charge', 
                      'velocity_magnitude', 'momentum_x', 'momentum_y', 'momentum_z', 
                      'theta_x', 'theta_y', 'f']
            
            data = []
            for i, splat in enumerate(splats):
                if len(splat) >= 10:  # 确保有足够的数据
                    x, y, z, mass, vx, vy, vz, ke, t, charge = splat[:10]
                    
                    # 获取扩展参数
                    if extended_params and i < len(extended_params.get('velocity_magnitude', [])):
                        velocity_magnitude = extended_params['velocity_magnitude'][i]
                        momentum_x = extended_params['momentum_x'][i]
                        momentum_y = extended_params['momentum_y'][i]
                        momentum_z = extended_params['momentum_z'][i]
                        theta_x = extended_params['theta_x'][i]
                        theta_y = extended_params['theta_y'][i]
                    else:
                        # 计算当前粒子的参数
                        velocity_magnitude = math.sqrt(vx**2 + vy**2 + vz**2)
                        momentum_x = mass * vx
                        momentum_y = mass * vy
                        momentum_z = mass * vz
                        theta_x = math.atan(vx / vz) if abs(vz) > 1e-10 else 0
                        theta_y = math.atan(vy / vz) if abs(vz) > 1e-10 else 0
                    
                    data.append([x, y, z, mass, vx, vy, vz, ke, t, charge,
                               velocity_magnitude, momentum_x, momentum_y, momentum_z,
                               theta_x, theta_y, frequency])
                else:
                    # 兼容旧格式
                    x, y, z, mass = splat[:4]
                    data.append([x, y, z, mass, 0, 0, 0, 0, 0, 0,
                               0, 0, 0, 0, 0, 0, frequency])
            
            # 写入CSV文件
            if CSVUtils.write_csv_data(target_path, data, headers, delimiter=','):
                return target_path
            else:
                return None
                
        except Exception as e:
            print(f"创建SIMION CSV文件失败: {e}")
            return None

    @staticmethod
    def create_simion_csv_with_emittance(splats, frequency, target_dir, emittance_data=None, prefix="simion"):
        """
        创建包含发射度信息的SIMION Splat数据CSV文件
        
        参数:
            splats: Splat数据列表 [(x, y, z, mass, vx, vy, vz, ke, t, charge), ...]
            frequency: 频率值
            target_dir: 目标目录
            emittance_data: 发射度数据字典
            prefix: 文件名前缀
            
        返回:
            str: 创建的文件路径，如果失败返回None
        """
        try:
            # 确保目标目录存在
            os.makedirs(target_dir, exist_ok=True)
            
            # 生成包含发射度信息的文件名（英文短标签，使用等号）
            if emittance_data:
                filename = (
                    f"freq={frequency:.0f}Hz_"
                    f"rms={emittance_data.get('rms_x', 0):.2e}_"
                    f"geom={emittance_data.get('geometric_x', 0):.2e}_"
                    f"norm={emittance_data.get('normalized_x', 0):.2e}.csv"
                )
            else:
                filename = f"freq={frequency:.0f}Hz.csv"
            target_path = os.path.join(target_dir, filename)
            
            # 准备数据
            headers = ['x', 'y', 'z', 'mass', 'vx', 'vy', 'vz', 'ke', 't', 'charge', 'f']
            data = []
            for splat in splats:
                if len(splat) >= 10:  # 确保有足够的数据
                    x, y, z, mass, vx, vy, vz, ke, t, charge = splat[:10]
                    data.append([x, y, z, mass, vx, vy, vz, ke, t, charge, frequency])
                else:
                    # 兼容旧格式
                    x, y, z, mass = splat[:4]
                    data.append([x, y, z, mass, 0, 0, 0, 0, 0, 0, frequency])
            
            # 写入CSV文件
            if CSVUtils.write_csv_data(target_path, data, headers, delimiter=','):
                return target_path
            else:
                return None
                
        except Exception as e:
            print(f"创建SIMION CSV文件失败: {e}")
            return None

    @staticmethod
    def create_simion_csv(splats, frequency, target_dir, prefix="simion"):
        """
        直接创建SIMION Splat数据的CSV文件
        
        参数:
            splats: Splat数据列表 [(x, y, z, mass), ...]
            frequency: 频率值
            target_dir: 目标目录
            prefix: 文件名前缀
            
        返回:
            str: 创建的文件路径，如果失败返回None
        """
        try:
            # 确保目标目录存在
            os.makedirs(target_dir, exist_ok=True)
            
            # 生成文件名
            filename = f"{frequency:.0f}.csv"
            target_path = os.path.join(target_dir, filename)
            
            # 准备数据
            headers = ['x', 'y', 'z', 'mass', 'vx', 'vy', 'vz', 'ke', 't', 'charge', 'f']
            data = []
            for splat in splats:
                if len(splat) >= 10:  # 确保有足够的数据
                    x, y, z, mass, vx, vy, vz, ke, t, charge = splat[:10]
                    data.append([x, y, z, mass, vx, vy, vz, ke, t, charge, frequency])
                else:
                    # 兼容旧格式
                    x, y, z, mass = splat[:4]
                    data.append([x, y, z, mass, 0, 0, 0, 0, 0, 0, frequency])
            
            # 写入CSV文件
            if CSVUtils.write_csv_data(target_path, data, headers, delimiter=','):
                return target_path
            else:
                return None
                
        except Exception as e:
            print(f"创建SIMION CSV文件失败: {e}")
            return None
    
    @staticmethod
    def read_csv_data(csv_path):
        """
        读取CSV文件数据
        
        参数:
            csv_path: CSV文件路径
            
        返回:
            list: CSV数据列表，如果失败返回空列表
        """
        try:
            data = []
            with open(csv_path, 'r', encoding='utf-8') as file:
                csv_reader = csv.reader(file)
                for row in csv_reader:
                    data.append(row)
            return data
        except Exception as e:
            print(f"读取CSV文件失败: {e}")
            return []
    
    @staticmethod
    def write_csv_data(csv_path, data, headers=None, delimiter=','):
        """
        写入CSV文件数据
        
        参数:
            csv_path: CSV文件路径
            data: 数据列表
            headers: 表头列表（可选）
            
        返回:
            bool: 写入成功返回True
        """
        try:
            # 确保目录存在
            os.makedirs(os.path.dirname(csv_path), exist_ok=True)
            
            with open(csv_path, 'w', newline='', encoding='utf-8') as file:
                csv_writer = csv.writer(file, delimiter=delimiter)
                
                # 写入表头
                if headers:
                    csv_writer.writerow(headers)
                
                # 写入数据
                for row in data:
                    csv_writer.writerow(row)
            
            return True
            
        except Exception as e:
            print(f"写入CSV文件失败: {e}")
            return False
    
    @staticmethod
    def format_csv_for_simion(data, frequency):
        """
        格式化CSV数据为SIMION格式
        
        参数:
            data: 原始数据
            frequency: 频率值
            
        返回:
            list: 格式化后的数据
        """
        try:
            formatted_data = []
            
            # 添加表头
            headers = ['x', 'y', 'z', 'mass', 'f']
            formatted_data.append(headers)
            
            # 处理数据行
            for row in data:
                if len(row) >= 5:
                    # 确保数据格式正确
                    formatted_row = [
                        float(row[0]) if row[0] else 0.0,
                        float(row[1]) if row[1] else 0.0,
                        float(row[2]) if row[2] else 0.0,
                        float(row[3]) if row[3] else 0.0,
                        float(row[4]) if row[4] else 0.0
                    ]
                    formatted_data.append(formatted_row)
            
            return formatted_data
            
        except Exception as e:
            print(f"格式化CSV数据失败: {e}")
            return []
    
    @staticmethod
    def get_csv_info(csv_path):
        """
        获取CSV文件信息
        
        参数:
            csv_path: CSV文件路径
            
        返回:
            dict: 文件信息字典
        """
        try:
            if not os.path.exists(csv_path):
                return {"exists": False}
            
            stat = os.stat(csv_path)
            data = CSVUtils.read_csv_data(csv_path)
            
            return {
                "exists": True,
                "size": stat.st_size,
                "rows": len(data),
                "modified": datetime.fromtimestamp(stat.st_mtime).strftime("%Y-%m-%d %H:%M:%S")
            }
            
        except Exception as e:
            print(f"获取CSV文件信息失败: {e}")
            return {"exists": False, "error": str(e)}


# 便捷函数
def create_simion_csv_file(splats, frequency, target_dir):
    """直接创建SIMION Splat数据的CSV文件"""
    return CSVUtils.create_simion_csv(splats, frequency, target_dir, "simion")


def create_frequency_csv(target_dir, frequency, data, emittance_data=None, extended_params=None):
    """
    为指定频率创建CSV文件，支持发射度和扩展参数
    
    参数:
        target_dir: 目标目录
        frequency: 频率值
        data: 数据列表
        emittance_data: 发射度数据字典（可选）
        extended_params: 扩展参数字典（可选）
    """
    # 生成包含所有信息的文件名（英文短标签，使用等号）
    if emittance_data and extended_params:
        filename = (
            f"freq={frequency:.0f}Hz_"
            f"rms={emittance_data.get('rms_x', 0):.2e}_"
            f"geom={emittance_data.get('geometric_x', 0):.2e}_"
            f"norm={emittance_data.get('normalized_x', 0):.2e}_"
            f"hull={extended_params.get('convex_hull_area', 0):.2e}.csv"
        )
    elif emittance_data:
        filename = (
            f"freq={frequency:.0f}Hz_"
            f"rms={emittance_data.get('rms_x', 0):.2e}_"
            f"geom={emittance_data.get('geometric_x', 0):.2e}_"
            f"norm={emittance_data.get('normalized_x', 0):.2e}.csv"
        )
    else:
        filename = f"freq={frequency:.0f}Hz.csv"
    
    target_path = os.path.join(target_dir, filename)
    
    return CSVUtils.write_csv_data(target_path, data, delimiter=',')
